import { makeLegalIdentifier } from '@rollup/pluginutils'
//@ts-ignore
import { traverse } from '@webassemblyjs/ast'
//@ts-ignore
import { decode } from '@webassemblyjs/wasm-parser'
import chalk from 'chalk'
import { readFile } from 'fs/promises'
import { PluginContext, PluginHooks } from 'rollup'
import {
  CustomWebAssemblyInstanceFactory,
  WebAssemblyModuleExportsTable,
  WebAssemblyModuleImportsTable,
} from '../types'
import { isWasmFromBuffer } from '../utils'
import { isWasmFromString } from './../utils'
import { getHelper as getBaseHelp } from './base'
import wasmLoaderHelper from './modules/wasm'
import { generageImports } from './modules/_utils'
import { defineHelper } from './_utils'

const decoderOpts = {
  ignoreCodeSection: true,
  ignoreDataSection: true,
}

const getWasmImportsExportsTable = (buffer: Buffer) => {
  const importsTable: WebAssemblyModuleImportsTable = Object.create(null)
  const exportsTable: WebAssemblyModuleExportsTable = Object.create(null)
  if (isWasmFromBuffer(buffer)) {
    const wasmAst = decode(buffer, decoderOpts)
    try {
      traverse(wasmAst, {
        ModuleImport({ node }: any) {
          const id = node.name
          const from = node.module
          let type = node.descr.type

          switch (type) {
            case 'FunctionType':
            case 'Function':
            case 'FuncType':
            case 'FuncImportDescr':
            case 'Func':
            case 'func':
              type = 'function'
              break
            case 'TableType':
            case 'TableImportDescr':
            case 'Table':
              type = 'table'
              break
            case 'GlobalImportDescr':
            case 'GlobalType':
            case 'Global':
              type = 'global'
              break
            case 'MemoryImportDescr':
            case 'MemoryType':
            case 'Memory':
              type = 'memory'
              break
          }

          ;(importsTable[from] ?? (importsTable[from] = Object.create(null)))[id] = type
        },
        ModuleExport({ node }: any) {
          const id = node.name
          let type = node.descr.exportType
          // "Func" | "Table" | "Memory" | "Global"
          switch (type) {
            case 'Func':
            case 'func':
              type = 'function'
              break
            case 'Table':
              type = 'table'
              break
            case 'Global':
              type = 'global'
              break
            case 'Memory':
              type = 'memory'
              break
          }
          exportsTable[id] = type
        },
      })
    } catch (error) {}
  }

  return { imports: importsTable, exports: exportsTable }
}

export const getHelper = defineHelper((options, registerVirtualModule) => {
  const utilsHelper = getBaseHelp(options)

  const _generateLoaderCode = async (
    context: PluginContext,
    id: string,
    fileName: string,
    loaderHelperModule = wasmLoaderHelper
  ) => {
    const source = await readFile(fileName)
    const factoryCode = options.useInstanceFactory?.(id, getWasmImportsExportsTable(source))
    const inline = options.inline(fileName, id)
    let factoryId = factoryCode ? registerVirtualModule(factoryCode.toString()) : ''

    if (inline) {
      return `
        ${generageImports(
          [loaderHelperModule, ['createInstantiateFromBuffer']],
          [utilsHelper.getInlineBufferModule(source), [], 'source']
        )}
        ${factoryId ? `import factory from ${JSON.stringify(factoryId)}` : ''};
        export const { instantiate, instantiateSync } = createInstantiateFromBuffer(source, ${
          factoryId ? 'factory' : ''
        });
        const _getExports = () => { try { return instantiateSync().exports } catch(err) { console.error(err) } };
        export default /*#__PURE__*/ _getExports();
      `
    } else {
      return `
        ${generageImports([loaderHelperModule, ['createInstantiateFromPath']])}
        ${factoryId ? `import factory from ${JSON.stringify(factoryId)}` : ''};
        const _fileName = ${await utilsHelper.fileToStringifyUrl(context, fileName, source)};
        export const { instantiate, instantiateSync } = 
          createInstantiateFromPath(_fileName, ${factoryId ? 'factory' : ''});
        const _getExports = () => { try { return instantiateSync().exports } catch(err) { console.error(err) } };
        export default /*#__PURE__*/ _getExports();
      `
    }
  }

  const load: PluginHooks['load'] = async function (id) {
    if (!id.endsWith('.wasm')) return
    try {
      utilsHelper.addWatchFile(this, id, id)

      return _generateLoaderCode(this, id, id)
    } catch (error) {
      console.error(chalk.red(error))
    }
  }

  return {
    _generateLoaderCode,
    load,
  }
})

export const defaultInstanceFactory: CustomWebAssemblyInstanceFactory = (id, { imports }) => {
  const importObjectStrs: string[] = []
  const importStatmentStrs: string[] = []

  let i = 0
  for (const [id, keys] of Object.entries(imports)) {
    const keyMapping = Object.keys(keys).map((key) => [
      key,
      makeLegalIdentifier(`$${i++}$${id}$${key}`),
    ])

    importStatmentStrs.push(
      `/*#__PURE__*/ import { ${keyMapping
        .map(([key, keyAs]) => `${key} as ${keyAs}`)
        .join(', ')} } from ${JSON.stringify(id)}`
    )
    importObjectStrs.push(
      `${JSON.stringify(id)}: { ${keyMapping
        .map(([key, keyAs]) => `${key}: ${keyAs}`)
        .join(', ')} }`
    )
  }

  return `${importStatmentStrs.join(';\n')};
    const defaultImportObject = /*#__PURE__*/ { ${importObjectStrs.join(',\n')} };
    export default (createInstance, __importObject) => {
      let importObject;
      if(__importObject) {
        importObject = { ...defaultImportObject }
        for(const modName in __importObject) {
          const mod = __importObject[modName];
          if(mod && typeof mod === 'object') {
            Object.assign(importObject[modName] = importObject[modName] ?? {}, mod);
          } else {
            importObject[modName] = mod;
          }
        }
      } else {
        importObject = defaultImportObject;
      }
      createInstance(importObject);
    }
  `
}
